#include <iostream>
#include <fstream>
#include <stdio.h>
#include <stdlib.h>
#include <cmath>
#include <time.h>

#include "opencv2/opencv.hpp"

using namespace cv;
using namespace std;

int find_median(int a[5])
{
	int swap;
	for (int c = 0 ; c < 5; c++)
	{
		for (int d = 0 ; d < 4- c ; d++)
			{
				 if (a[d] > a[d+1]) /* For decreasing order use < */
				 {
					  swap   = a[d];
					  a[d]   = a[d+1];
					  a[d+1] = swap;
				}
			}
	}
	return(a[2]);
}


int main(int, char** argv)
{
    VideoCapture cap(argv[1]); // open the file
    if(!cap.isOpened())  // check if we succeeded
        return -1;

    /////////////////////////////////////// Variable Declaration /////////////////////////////////////////////////////

    Mat frame;
    Mat edges;
    Mat prev_frame;
    Mat cur_frame;
    Mat next_frame;
    Size image_size((int)cap.get(CV_CAP_PROP_FRAME_WIDTH),(int)cap.get(CV_CAP_PROP_FRAME_HEIGHT));
    Mat cur_temp=Mat::zeros(image_size,CV_8UC1);
    Mat cur_temp_smooth=Mat::zeros(image_size,CV_8UC1);
    Mat prev_temp=Mat::zeros(image_size,CV_8UC1);    
    Mat cur_edge=Mat::zeros(image_size,CV_8UC1);
    Mat prev_edge=Mat::zeros(image_size,CV_8UC1);
    Mat next_edge=Mat::zeros(image_size,CV_8UC1);
	Mat next_temp=Mat::zeros(image_size,CV_8UC1);
	Mat prev_temp2=Mat::zeros(image_size,CV_8UC1);
	Mat next_temp2=Mat::zeros(image_size,CV_8UC1);
    Mat img_diff;
    Mat diff3 = Mat::zeros(image_size,CV_8UC1);
    Mat gradx = Mat::zeros(image_size,CV_8UC1); 
	Mat grady = Mat::zeros(image_size,CV_8UC1); 
	Mat grad = Mat::zeros(image_size,CV_8UC1); 
    Mat mask = Mat::zeros(image_size,CV_8UC1); 
	Mat background = Mat::zeros(image_size,CV_8UC1);
    Mat diff_staticedge = Mat::zeros(image_size,CV_8UC1); 
    Mat diff_moveedge = Mat::zeros(image_size,CV_8UC1); 
    Mat recon_image = Mat::zeros(image_size,CV_8UC1); 
    Mat range = Mat::zeros(image_size,CV_8UC1);
    Mat spreadsymm = Mat::zeros(image_size,CV_8UC1);
    Mat classify = Mat::zeros(image_size,CV_8UC1);

    vector<float> tpf,tnf;
    int max,min;
    float mean,T1,T2,below_std,above_std;
	clock_t t1,t2;
    float diff=0,time_frame=0;
	char* cthreshold=argv[2];
	char* cratio=argv[3];
	float threshold=atof(cthreshold);
	float ratio=atof(cratio);
	int no_frames=0;
    T1=40;
    T2=4;

    double fps=cap.get(CV_CAP_PROP_FPS);
    VideoWriter writer("output_video.avi",CV_FOURCC('M', 'J', 'P', 'G'),fps, image_size);


    /////////////////////////////////////////// File Opening and setting /////////////////////////////////////////////////////

	ofstream res_file;

    res_file.open("file.txt");

	if(res_file.is_open())
	{
        cout<<"File open.\n";
        res_file<<"Results:\n";
	}
	else	
		cout<<"Unable to open the file.\n";

    ///////////////////////******************************** Code **********************************/////////////////////////

    for( ; ; )
    {
		no_frames++;
        cap>>frame; 
        //imshow("prev_video", prev_temp);        
        imshow("current video", cur_temp);        
        cvtColor(frame, next_temp2, CV_BGR2GRAY);
        //imshow("next_video", next_temp);

		t1=clock();
                                  
  //////////////////////////////////////////////////  candidate rain pixels ////////////////////////////////////////////////
        
        uchar* data_cur=(uchar*)cur_temp.data;
        uchar* data_prev=(uchar*)prev_temp.data;
		uchar* data_prev2=(uchar*)prev_temp2.data;
        uchar* data_next=(uchar*)next_temp.data;
		uchar* data_next2=(uchar*)next_temp2.data;
        uchar* data_diff3=(uchar*)diff3.data;
		uchar* data_grad=(uchar*)grad.data;
		int width_step=cur_temp.step1();

        for(int i=0;i<image_size.height;i++)
        {
			for(int j=0; (j<((0.2)*image_size.width));j++)
           {
             int k_cur = data_cur[i*width_step+j];
             int k_pre = data_prev[i*width_step+j];
             int k_next = data_next[i*width_step+j];             
             if(((k_cur-k_pre)>=7) || ((k_cur-k_next)>=7))//|| (((2*k_cur-k_pre-k_next)==0)))
              {
                data_diff3[i*width_step+j]=255;
              } 
             else
                data_diff3[i*width_step+j]=0;                           
           }
         }

		for(int i=0;i<image_size.height;i++)
        {
			for(int j=(image_size.width-1);j>=((0.8)*image_size.width);j--)
           {
             int k_cur = data_cur[i*width_step+j];
             int k_pre = data_prev[i*width_step+j];
             int k_next = data_next[i*width_step+j];             
             if(((k_cur-k_pre)>=7) || ((k_cur-k_next)>=7))//|| (((2*k_cur-k_pre-k_next)==0)))
              {
                data_diff3[i*width_step+j]=255;
              } 
             else
                data_diff3[i*width_step+j]=0;                           
           }
         }
        
		for(int i=0;i<(0.2)*image_size.height;i++)
        {
			for(int j=(0.2*image_size.width);j<(0.8*image_size.width);j++)
           {
             int k_cur = data_cur[i*width_step+j];
             int k_pre = data_prev[i*width_step+j];
             int k_next = data_next[i*width_step+j];             
             if(((k_cur-k_pre)>=7) || ((k_cur-k_next)>=7))//|| (((2*k_cur-k_pre-k_next)==0)))
              {
                data_diff3[i*width_step+j]=255;
              } 
             else
                data_diff3[i*width_step+j]=0;                           
           }
        }

        //imshow("diff", diff3);       // All rain pixels
        
    //////////////////////////////////////////// static edge detection and removal /////////////////////////////////////////
    
        boxFilter(cur_temp,cur_temp_smooth,8, Size(3,3)); // changed from 5,5
        Canny(prev_temp,prev_edge,threshold,ratio*threshold,3);              // Threshold =10 Ratio =12.5
        Canny(next_temp,next_edge,threshold,ratio*threshold,3);
        Canny(cur_temp,cur_edge,threshold,ratio*threshold,3);

        uchar* data_prev_edge=(uchar*)prev_edge.data;
        uchar* data_next_edge=(uchar*)next_edge.data;
        uchar* data_cur_edge=(uchar*)cur_edge.data;
		uchar* data_diff_staticedge=(uchar*)diff_staticedge.data;
        uchar* data_edge=(uchar*)edges.data;
        
        Canny(cur_temp_smooth, edges,threshold,ratio*threshold,3); 

		//imshow("Edges",edges);
        
        for(int i=0;i<image_size.height;i++)
        {
			for(int j=0; (j<((0.2)*image_size.width));j++)
           {
           if((data_prev_edge[i*width_step+j]==255)&&(data_cur_edge[i*width_step+j]==255)&&(data_next_edge[i*width_step+j]==255))            
              data_diff_staticedge[i*width_step+j]=255;  
           else
              data_diff_staticedge[i*width_step+j]=0;               
           }
        } 

        for(int i=0;i<image_size.height;i++)
        {
			for(int j=(image_size.width-1);j>=((0.8)*image_size.width);j--)
           {
           if((data_prev_edge[i*width_step+j]==255)&&(data_cur_edge[i*width_step+j]==255)&&(data_next_edge[i*width_step+j]==255))            
              data_diff_staticedge[i*width_step+j]=255;  
           else
              data_diff_staticedge[i*width_step+j]=0;               
           }
        } 

		for(int i=0;i<(0.2)*image_size.height;i++)
        {
			for(int j=(0.2*image_size.width);j<(0.8*image_size.width);j++)
           {
           if((data_prev_edge[i*width_step+j]==255)&&(data_cur_edge[i*width_step+j]==255)&&(data_next_edge[i*width_step+j]==255))            
              data_diff_staticedge[i*width_step+j]=255;  
           else
              data_diff_staticedge[i*width_step+j]=0;               
           }
        } 
        
        for(int i=0;i<image_size.height;i++)
        {
          for(int j=0;j<image_size.width;j++)
           {              
                data_diff3[i*width_step+j]=data_diff3[i*width_step+j]-data_diff_staticedge[i*width_step+j];                                       
           }
        }
        
        std::vector<std::vector<Point> > contour;
        findContours(diff3,contour,CV_RETR_LIST,CV_CHAIN_APPROX_NONE);
        Mat cont = Mat::zeros(image_size,CV_8UC3);
        std::vector<std::vector<Point> > contour_boundary;
        
        for(int k=0;k<contour.size();k++)
        {
          double area = contourArea(contour[k]);
          if(area>=8 && area<=40)
          {
            contour_boundary.push_back(contour[k]);            
          }
        } 
        
        drawContours(cont,contour_boundary,-1,Scalar(255,255,255),-1);
        
        std::vector<Mat> cont_filled;
        split(cont,cont_filled);        
        
        diff3=diff3-cont_filled[0];
		Mat element = getStructuringElement( 2, Size( 3, 3 ), Point( 1, 1 ) );

        /// Apply the dilation and erosion
        dilate( diff3, diff3, element );
		erode (diff3,diff3,element);

		//Apply scharr operator in x and y direction

		Scharr( diff3, gradx, diff3.depth(), 1, 0, 1, 0, BORDER_DEFAULT );
		convertScaleAbs( gradx,gradx );
		Scharr( diff3, grady, diff3.depth(), 0, 1, 1, 0, BORDER_DEFAULT );
		convertScaleAbs( grady,grady);
		
		addWeighted( gradx, 0.5, grady, 0.5, 0, grad );
        
		//imshow("diff_refined",diff3);        
        //imshow("contour",cont);         

        
   ////////////////////////////////////////////////// moving edge detection ////////////////////////////////////////////     
        
        cur_edge.copyTo(diff_moveedge);

        uchar* data_diff_moveedge=(uchar*)diff_moveedge.data;
        
        for(int i=0;i<image_size.height;i++)
        {
			for(int j=0; (j<((0.2)*image_size.width));j++)
            {
                 int temp=data_cur_edge[i*width_step+j];
                 int temp_pre=data_prev_edge[i*width_step+j];
                 int temp_next=data_next_edge[i*width_step+j];
                 if ((temp==temp_pre) && (temp==temp_next))
                 {
                    data_diff_moveedge[i*width_step+j]=0;
                 }
            }
        }

		for(int i=0;i<image_size.height;i++)
        {
			for(int j=(image_size.width-1);j>=((0.8)*image_size.width);j--)
           {
                 int temp=data_cur_edge[i*width_step+j];
                 int temp_pre=data_prev_edge[i*width_step+j];
                 int temp_next=data_next_edge[i*width_step+j];
                 if ((temp==temp_pre) && (temp==temp_next))
                 {
                    data_diff_moveedge[i*width_step+j]=0;
                 }
            }
        }

		for(int i=0;i<(0.2)*image_size.height;i++)
        {
			for(int j=(0.2*image_size.width);j<(0.8*image_size.width);j++)
           {
                 int temp=data_cur_edge[i*width_step+j];
                 int temp_pre=data_prev_edge[i*width_step+j];
                 int temp_next=data_next_edge[i*width_step+j];
                 if ((temp==temp_pre) && (temp==temp_next))
                 {
                    data_diff_moveedge[i*width_step+j]=0;
                 }
            }
        }
        
		dilate( diff_moveedge, diff_moveedge, element );
		//erode(diff_moveedge,diff_moveedge,element);

        diff3=diff3-diff_moveedge; 
       
		//diff4=next_temp-diff3;
                 
        //imshow("diff_moveedge",diff_moveedge);
        //imshow("Only edges",diff3);

		//Edges 
		//imshow("x edges",gradx);
		//imshow("y edges",grady);
		//imshow("edges after scharr",grad);

        grad=grad-diff_moveedge;

		//imshow("edges after removing moving edge",grad);

   //////////////////////////////////////////////////// Feature Extraction ////////////////////////////////////////////////
  
  /**************************************** Range ****************************************/
   uchar* data_range=(uchar*)range.data;
   
   for(int i=0;i<image_size.height;i++)
   {
    for(int j=0;j<image_size.width;j++)
     {
       max= data_cur[i*width_step+j];
       min= data_cur[i*width_step+j];
       if(max < data_next[i*width_step+j] )
       {
        max = data_next[i*width_step+j];
       }
       else if ( max < data_prev[i*width_step+j])
       {
        max = data_prev[i*width_step+j];
       }
       if (min > data_next[i*width_step+j])
       {
        min = data_next[i*width_step+j];
       }
       else if (min > data_prev[i*width_step+j])
       {
        data_next[i*width_step+j];
       }
       
       data_range[i*width_step+j]= max-min;        
     }
   }

   /********************************************* Spread Symmetry ***************************************/

  uchar* data_spreadsymm = (uchar*)spreadsymm.data;
   
   for(int i=0;i<image_size.height;i++)
   {
    for(int j=0;j<image_size.width;j++)
     {
       mean = (data_cur[i*width_step+j]+data_next[i*width_step+j]+data_cur[i*width_step+j])/3.0;

       if(data_prev[i*width_step+j]<mean)
       {
           if(data_cur[i*width_step+j]<mean)
           {
               above_std=0;
               below_std=(float)sqrt(pow((data_prev[i*width_step+j]-mean),2)+pow((data_cur[i*width_step+j]-mean),2))/2.0;
           }
       }

       else if (data_cur[i*width_step+j]>mean)
       {
           if(data_next[i*width_step+j]>mean)
           {
               below_std=0;
               above_std=(float)sqrt(pow((data_cur[i*width_step+j]-mean),2)+pow((data_next[i*width_step+j]-mean),2))/2.0;
           }
       }

       else if(data_cur[i*width_step+j]<mean)
       {
           if(data_next[i*width_step+j]>mean)
           {
               below_std=0;
               above_std=(float)sqrt(pow((data_prev[i*width_step+j]-mean),2)+pow((data_next[i*width_step+j]-mean),2))/2.0;
           }
       }

       else if(data_next[i*width_step+j]<mean)
       {
           if(data_prev[i*width_step+j]>mean)
           {
               below_std=0;
               above_std=(float)sqrt(pow((data_prev[i*width_step+j]-mean),2)+pow((data_cur[i*width_step+j]-mean),2))/2.0;
           }
       }
       else if(data_cur[i*width_step+j]<mean)
       {
           if(data_next[i*width_step+j]<mean)
           {
               above_std=0;
               below_std=(float)sqrt(pow((data_cur[i*width_step+j]-mean),2)+pow((data_next[i*width_step+j]-mean),2))/2.0;
           }
       }
       else if(data_next[i*width_step+j]<mean)
       {
           if(data_prev[i*width_step+j]<mean)
           {
               above_std=0;
               below_std=(float)sqrt(pow((data_prev[i*width_step+j]-mean),2)+pow((data_next[i*width_step+j]-mean),2))/2.0;
           }
       }
       data_spreadsymm[i*width_step+j]=abs(above_std-below_std);
     }
   }  

   //imshow("Range",range);
   //imshow("Spreadsymmetry",spreadsymm);

   //////////////////////////////////////////////////// Temporal Median //////////////////////////////////////////////////

   uchar* data_background = (uchar*)background.data;
   int a[5],img_median;

   for(int i=0;i<image_size.height;i++)
   {
    for(int j=0;j<image_size.width;j++)
    {
		a[0]=data_prev2[i*width_step+j];
		a[1]=data_prev[i*width_step+j];
		a[2]=data_cur[i*width_step+j];
		a[3]=data_next[i*width_step+j];
		a[4]=data_next2[i*width_step+j];

		img_median=find_median(a);
		data_background[i*width_step+j]=img_median;
	}
   }

   //imshow("Background",background);

   //////////////////////////////////////////////////// Classification ////////////////////////////////////////////////////

   uchar* data_classify = (uchar*)classify.data;
   
   /*for(int i=0;i<image_size.height;i++)
   {
    for(int j=0;j<image_size.width;j++)
    {
      if( (data_range[i*width_step+j]>8) && (data_spreadsymm[i*width_step+j]>10))
        data_classify[i*width_step+j]=255;

	    data_classify[i*width_step+j]=255-data_classify[i*width_step+j];
    }    
   }*/

   for(int i=0;i<image_size.height;i++)
        {
		for(int j=0; (j<((0.2)*image_size.width));j++)
            {
				if( (data_range[i*width_step+j]>8) && (data_spreadsymm[i*width_step+j]>10))
				data_classify[i*width_step+j]=255;

				data_classify[i*width_step+j]=255-data_classify[i*width_step+j];
			}    
		}

   for(int i=0;i<image_size.height;i++)
        {
			for(int j=(image_size.width-1);j>=((0.8)*image_size.width);j--)
           {
				if( (data_range[i*width_step+j]>8) && (data_spreadsymm[i*width_step+j]>10))
				data_classify[i*width_step+j]=255;

				data_classify[i*width_step+j]=255-data_classify[i*width_step+j];
		   }    
        }

   for(int i=0;i<(0.2)*image_size.height;i++)
        {
			for(int j=(0.2*image_size.width);j<(0.8*image_size.width);j++)
           {
			if( (data_range[i*width_step+j]>8) && (data_spreadsymm[i*width_step+j]>10))
			data_classify[i*width_step+j]=255;

			data_classify[i*width_step+j]=255-data_classify[i*width_step+j];
		   }    
		}

   imshow("Classifier",classify);

   //////////////////////////////////////////////////// Parameter Calculation /////////////////////////////////////////////
   
   /*********************************************** True Positives and False Negatives ************************************/
   /*
   int count_tp,int count tn;
   
   for(int i=0;i<image_size.height;i++)
   {
    for(int j=0;j<image_size.width;j++)
    {
      if( (data_classify[i*width_step+j]==) && ())
      { 
        count_tp++; 
        tpf.push_back(count_tp++);
      }
      else 
      {
      count_tn++;
        tnf_pushback(count_tn);
      }
    }    
   }    
   
   */
   //////////////////////////////////////////////////// Image Inpainting /////////////////////////////////////////////////
        
        uchar* data_reconimg=(uchar*)recon_image.data;

		/*for(int p=0;p<image_size.height;p++)
        {
			for(int q=0;q<image_size.width;q++)
            {
            //if((data_grad[p*width_step+q])||data_classify[p*width_step+q])
			if(data_classify[p*width_step+q])
              { 
                data_reconimg[p*width_step+q]=(data_prev[p*width_step+q]+data_next[p*width_step+q]+data_prev2[p*width_step+q]+data_next2[p*width_step+q])/4;                
              }
              else
              {
                data_reconimg[p*width_step+q]=data_cur[p*width_step+q];
              }               
            }
        }*/
        
        for(int p=0;p<image_size.height;p++)
        {
			for(int q=0; (q<((0.2)*image_size.width));q++)
            {
            //if((data_grad[p*width_step+q])||data_classify[p*width_step+q])
			if(data_classify[p*width_step+q])
              { 
                data_reconimg[p*width_step+q]=(data_prev[p*width_step+q]+data_next[p*width_step+q]+data_prev2[p*width_step+q]+data_next2[p*width_step+q])/4;                
              }
              else
              {
                data_reconimg[p*width_step+q]=data_cur[p*width_step+q];
              }               
            }
        }

		for(int p=0;p<image_size.height;p++)
        {
			for(int q=(image_size.width-1);q>=((0.8)*image_size.width);q--)
           {
            //if((data_grad[p*width_step+q])||data_classify[p*width_step+q])
			if(data_classify[p*width_step+q])
              { 
                data_reconimg[p*width_step+q]=(data_prev[p*width_step+q]+data_next[p*width_step+q]+data_prev2[p*width_step+q]+data_next2[p*width_step+q])/4;                
              }
              else
              {
                data_reconimg[p*width_step+q]=data_cur[p*width_step+q];
              }               
            }
        }

		for(int p=0;p<(0.2)*image_size.height;p++)
        {
			for(int q=(0.2*image_size.width);q<(0.8*image_size.width);q++)
           {
            //if((data_grad[p*width_step+q])||data_classify[p*width_step+q])
			if(data_classify[p*width_step+q])
              { 
                data_reconimg[p*width_step+q]=(data_prev[p*width_step+q]+data_next[p*width_step+q]+data_prev2[p*width_step+q]+data_next2[p*width_step+q])/4;                
              }
              else
              {
                data_reconimg[p*width_step+q]=data_cur[p*width_step+q];
              }               
            }
        }

       // GaussianBlur(recon_image,recon_image,Size(1,1),1.5,1.5);

  ////////////////////////////////////////////// Writing time in file ////////////////////////////////////////////////////
		t2=clock();
        diff=float(t2-t1)/CLOCKS_PER_SEC;
        time_frame+= diff;
        res_file<<"Time for frame "<<no_frames<<" is: "<<diff<<" s."<<endl;

        
   ////////////////////////////////////////////////// writing video file /////////////////////////////////////////////////

		imshow("Reconstructed Image",recon_image);
        Mat color = Mat::zeros(image_size,CV_8UC3);
        Mat gray[] = {recon_image,recon_image,recon_image};
        merge(gray,3,color);
        writer.write(recon_image);
        
   ///////////////////////////////////////////////// Taking next frame ///////////////////////////////////////////////////             

		prev_temp.copyTo(prev_temp2);
        cur_temp.copyTo(prev_temp);
        next_temp.copyTo(cur_temp);
		next_temp2.copyTo(next_temp);
                      
        if(waitKey(150) >= 0) break;
    }

    cap.release();

    res_file<<"Average Time per frame:"<<time_frame/no_frames<<endl;
    res_file.close();
	
    return 0;
}
